import simplejson
from redis import StrictRedis, ConnectionPool
from util.const import REDIS_URL, WS_DEPTH, WS_TRADE
from functools import partial
from util.jiafu_log import JiaFuLogger
from util.jiafu_ws_api import WsApi
from util.jiafu_rest_api import RestApi
from util.const import *


logger = JiaFuLogger("jiafu_server:utils:common_func")

POOL = ConnectionPool.from_url(REDIS_URL)


def singleton(cls):
    instances = {}

    def _singleton(*args, **kwargs):
        if cls not in instances:
            instances[cls] = cls(*args, **kwargs)
        return instances[cls]

    return _singleton


def get_redis_con():
    return StrictRedis(connection_pool=POOL)


class SubSymbol(object):
    def __init__(self, symbol):
        self.symbol = symbol
        self.sub_symbol(symbol)

    def sub_symbol(self, symbol):
        con = get_redis_con()
        self.ps = con.pubsub()
        self.ps.subscribe(symbol + '_DEPTH', symbol + '_TRADE')

    def get_message(self, timeout=30):
        """
        获取币对的深度和交易推送，该方法阻塞直到有收到数据
        :param timeout: 阻塞时间，不能设置的太短
        :return:
        """
        msg = self.ps.get_message(timeout=timeout)
        ws_need_reconnect = False
        ret = None
        logger.debug(f'get redis msg {msg}')
        if not msg:
            ws_need_reconnect = True
        if msg and msg["type"] == "message":
            if msg["data"] != b"still alive":
                try:
                    data = simplejson.loads(msg["data"])
                except Exception as e:
                    logger.error('ws msg json loads error msg:  %s, data:%s' % (str(e), msg))
                else:
                    if msg["channel"].endswith(b"TRADE"):
                        ret = {"type": "trade", "msg": data["message"]}
                    elif msg["channel"].endswith(b"DEPTH"):
                        ret = {"type": "depth", "msg": data["message"]}
        return ws_need_reconnect, ret

@singleton
class PubSymbol(object):
    NEED_RECONNECT = True

    def __init__(self, url=None, cookies=None):
        self.sub_symbols_depth = {}
        self.sub_symbols_trade = {}
        self.con = get_redis_con()
        if url:
            self.url = url
        else:
            self.url = GZ_WS_BASE_URL
        if cookies:
            self.cookie = cookies
        else:
            self.cookie = self.get_ws_cookies()

    def get_ws_cookies(self):
        restApi = RestApi(GZ_REST_BASE_URL)
        res = restApi.login(name=WS_USER, password=WS_USER_PWD)
        if res['errno'] == 0:
            return restApi.get_cookies()

    def set_symbol_ws(self, symbol):
        if symbol not in self.sub_symbols_depth:
            ws = WsApi(self.url, self.cookie, WS_DEPTH)
            self.sub_symbols_depth[symbol] = ws
            ws.run(symbol=symbol,  message_recv_callback=partial(self.__pub_symbol, symbol + "_DEPTH"))

        if symbol not in self.sub_symbols_trade:
            ws = WsApi(self.url, self.cookie, WS_TRADE)
            self.sub_symbols_trade[symbol] = ws
            ws.run(symbol=symbol,  message_recv_callback=partial(self.__pub_symbol, symbol + "_TRADE"))

    def stop_pub(self, symbol=None):
        if symbol:
            if symbol in self.sub_symbols_depth:
                ws = self.sub_symbols_depth[symbol]
                ws.close()
            if symbol in self.sub_symbols_trade:
                ws = self.sub_symbols_trade[symbol]
                ws.close()
        else:
            for symbol in self.sub_symbols_depth:
                ws = self.sub_symbols_depth[symbol]
                ws.close()
            for symbol in self.sub_symbols_trade:
                ws = self.sub_symbols_trade[symbol]
                ws.close()

    def reconnect(self):
        """
        目前发现ws连接一段时间后，没有收到对端ws服务器的推送，这里重新连接一下ws服务器，看下能否收到推送
        :return:
        """
        if not self.NEED_RECONNECT:
            import time
            logger.debug('wait ws reconnect')
            time.sleep(3)
            return
        else:
            self.NEED_RECONNECT = False
            logger.info('reconnect to binance ws.....')
            # 停止所有ws连接
            self.stop_pub()
            # 更新下cookies
            self.cookie = self.get_ws_cookies()
            for symbol in self.sub_symbols_trade:
                ws_depth = WsApi(self.url, self.cookie, WS_DEPTH)
                self.sub_symbols_depth[symbol] = ws_depth
                ws_depth.run(symbol=symbol, message_recv_callback=partial(self.__pub_symbol, symbol + "_DEPTH"))
                ws_trade = WsApi(self.url, self.cookie, WS_TRADE)
                self.sub_symbols_trade[symbol] = ws_trade
                ws_trade.run(symbol=symbol, message_recv_callback=partial(self.__pub_symbol, symbol + "_TRADE"))
            self.NEED_RECONNECT = True

    def __pub_symbol(self, channel, msg):
        self.con.publish(channel, msg)






